# coach/views.py
from django.shortcuts import render
from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework import status
from langchain.chains import LLMChain, RetrievalQA
from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
from langchain.memory import ConversationBufferMemory

from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_openai import ChatOpenAI

from langchain.schema import Document
from langchain.chains import ConversationalRetrievalChain
from django.apps import apps
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

# 获取API密钥
DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")

# 初始化Qwen模型
llm = ChatOpenAI(
    model="qwen-plus",
    api_key=DASHSCOPE_API_KEY,
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0.7,
    max_tokens=2048
)

# 初始化阿里云Embeddings服务
embeddings = DashScopeEmbeddings(
    dashscope_api_key=DASHSCOPE_API_KEY,
    model="text-embedding-v1"  # 使用阿里云文本嵌入模型
)

# 创建对话模板
prompt = ChatPromptTemplate.from_messages([
    HumanMessagePromptTemplate.from_template("""你现在是一名专业的AI健身教练，擅长提供科学的健身指导、营养建议和健康管理方案。
    你的回答需要基于用户的具体情况，提供可操作的建议。

    【回答要求】
    1. 优先使用检索到的健康知识库内容进行回答，确保回答的专业性和准确性
    2. 如果检索结果不充分，请结合你的专业知识提供个性化建议
    3. 对于健身相关问题，需考虑用户可能的运动水平（初学者/中级/高级）
    4. 提供具体的行动步骤或训练计划时，需明确指出适用人群
    5. 遇到健康风险相关问题，建议用户咨询专业人士

    【检索结果】
    {context}

    【历史对话】
    {chat_history}

    用户: {question}
    Coach:""")
])

# 初始化内存
memory = ConversationBufferMemory(
    memory_key="chat_history",
    input_key="question",
    output_key="answer",
    return_messages=True
)


# 加载健康表数据并创建向量索引
def load_health_data():
    try:
        HealthFAQ = apps.get_model('coach', 'HealthFAQ')  # 替换为实际应用名
        documents = []

        for faq in HealthFAQ.objects.all():
            content = f"问题：{faq.question}\n回答：{faq.answer}"
            documents.append(Document(
                page_content=content,
                metadata={"question": faq.question, "answer": faq.answer}
            ))

        # 创建向量索引
        if documents:
            vectorstore = FAISS.from_documents(documents, embeddings)
            return vectorstore
        return None
    except Exception as e:
        print(f"加载健康数据时出错: {e}")
        return None


# 创建检索链的函数
def create_coach_chain():
    vectorstore = load_health_data()
    
    if vectorstore:
        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

        # 创建问题转换提示词
        condense_question_prompt = ChatPromptTemplate.from_messages([
            HumanMessagePromptTemplate.from_template("""根据历史对话和当前问题，生成一个独立的问题。
            历史对话: {chat_history}
            当前问题: {question}
            独立问题:""")
        ])

        # 创建对话检索链
        return ConversationalRetrievalChain.from_llm(
            llm=llm,
            retriever=retriever,
            condense_question_prompt=condense_question_prompt,
            chain_type="stuff",
            return_source_documents=True,
            memory=memory,
            get_chat_history=lambda h: h
        ), vectorstore
    else:
        # 如果没有数据，回退到普通LLM链
        return LLMChain(
            llm=llm,
            prompt=prompt,
            memory=memory,
            verbose=True
        ), None


class CoachAPIView(APIView):
    """Coach AI助手API"""

    def post(self, request, *args, **kwargs):
        user_input = request.data.get('user_input')

        if not user_input:
            return Response(
                {"error": "user_input字段是必需的"},
                status=status.HTTP_400_BAD_REQUEST
            )

        try:
            # 在视图函数内部创建链
            coach_chain, vectorstore = create_coach_chain()
            
            # 调用LangChain链获取响应
            if vectorstore:
                # 获取当前历史对话
                chat_history = memory.load_memory_variables({})['chat_history']

                # 准备输入参数
                inputs = {
                    "question": user_input,
                    "chat_history": chat_history
                }

                # 调用对话检索链
                result = coach_chain(inputs)
                response = result['answer']
                source_docs = result.get('source_documents', [])

                # 简化来源信息
                sources = []
                for doc in source_docs:
                    sources.append({
                        "question": doc.metadata.get("question"),
                        "answer": doc.metadata.get("answer"),
                        "relevance_score": doc.metadata.get("score", "N/A")  # 添加相关性分数
                    })

                # 简化历史对话格式
                simplified_history = []
                for msg in chat_history:
                    if msg.type == "human":
                        simplified_history.append({"user": msg.content})
                    elif msg.type == "ai":
                        simplified_history.append({"coach": msg.content})

                return Response({
                    "user_input": user_input,
                    "response": response,
                    "sources": sources,
                    "conversation_history": simplified_history
                })
            else:
                # 没有向量索引时的处理
                response = coach_chain.run(user_input=user_input)

                # 获取简化的历史对话
                chat_history = memory.load_memory_variables({})['chat_history']
                simplified_history = []
                for msg in chat_history:
                    if msg.type == "human":
                        simplified_history.append({"user": msg.content})
                    elif msg.type == "ai":
                        simplified_history.append({"coach": msg.content})

                return Response({
                    "user_input": user_input,
                    "response": response,
                    "conversation_history": simplified_history
                })

        except Exception as e:
            return Response(
                {"error": f"处理请求时发生错误: {str(e)}"},
                status=status.HTTP_500_INTERNAL_SERVER_ERROR
            )

